import sys, os
import cv2
import argparse
import numpy as np
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PIL import Image, ImageEnhance
import label_eval


class ImageLabel(QLabel):
    croped = pyqtSignal()
    resized = pyqtSignal()
    def __init__(self, parent=None, mode='text'):
        super(ImageLabel, self).__init__()
        self.parent = parent
        self.isCrop = False
        self.image = None
        self.cropImage = None
        self.affine = None
        self.bgn, self.end, self.curr = None, None, None
        self.imgW, self.imgH = 0, 0
        self.rawImgW, self.rawImgH = 0, 0
        self.cropX1 = 0
        self.cropX2 = 0
        self.cropY1 = 0
        self.cropY2 = 0
        self.mode = mode
        #self.test = None
        #self.test2 = None

    def showImage(self, imgpath):
        self.image = cv2.imread(imgpath)
        img = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
        img = QImage(img[:], img.shape[1], img.shape[0], img.shape[1] * 3, QImage.Format_RGB888)
        self.rawImgW, self.rawImgH = img.size().width(), img.size().height()
        img = img.scaled(QSize(self.width(), self.height()), Qt.KeepAspectRatio)
        self.imgW, self.imgH = img.width(), img.height()
        self.setPixmap(QPixmap.fromImage(img))
        #print('lbSize: {}, imgSize: {}, rawImgSize: {}'.format((self.width(), self.height()), (self.imgW, self.imgH), (self.rawImgW, self.rawImgH)))

    def setAffine(self, ofs):
        if self.cropImage is not None:
            w, h = self.cropImage.shape[1], self.cropImage.shape[0]
            ##src-->dst(左上角，左下角，右上角）
            if ofs > 0:
                src=np.float32([[0,0],[0+ofs,h-1],[w-1,0]])
                dst=np.float32([[0,0],[0,h-1],[w-1,0]])
            else:
                src=np.float32([[0-ofs,0],[0,h-1],[w-1,0]])
                dst=np.float32([[0,0],[0,h-1],[w-1+ofs,0]])
            #cv2.getPerspectiveTransform -- 4点仿射
            self.affine = cv2.getAffineTransform(src, dst)
            print(self.affine)

    def crop(self):
        rectX, rectY = self.bgn[0], self.bgn[1]
        rectW, rectH = self.end[0] - self.bgn[0], self.end[1] - self.bgn[1]
        if rectW * rectH > 100:
            imgX, imgY = (self.width() - self.imgW)//2, (self.height() - self.imgH)//2
            cropX1 = self.rawImgW * (rectX - imgX) // self.imgW
            cropY1 = self.rawImgH * (rectY - imgY) // self.imgH
            cropX2 = self.rawImgW * (rectX - imgX + rectW) // self.imgW
            cropY2 = self.rawImgH * (rectY - imgY + rectH) // self.imgH
            print('rect: ({}, {}, {}, {})'.format(rectX, rectY, rectW, rectH))
            print('img : ({}, {}, {}, {})'.format(imgX, imgY, self.imgW, self.imgH))
            print('crop: ({}, {}, {}, {})'.format(cropX1, cropY1, cropX2-cropX1, cropY2-cropY1))
            #cv2.imshow('crop', self.image[cropY1:cropY2, cropX1:cropX2])
            #cv2.waitKey(0)
            self.cropImage = cv2.cvtColor(self.image[cropY1:cropY2, cropX1:cropX2], cv2.COLOR_BGR2RGB)
            self.cropX1 = cropX1
            self.cropX2 = cropX2
            self.cropY1 = cropY1
            self.cropY2 = cropY2
            
            print(self.mode)
            if self.mode == 'mask':
                import snapmodel
                img = self.image[cropY1:cropY2, cropX1:cropX2].copy()
                #img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
                #img = self.cropImage.copy()
                cv2.imshow('src', img)
                bg = cv2.imread('bg.jpg')[self.cropY1:self.cropY2, self.cropX1:self.cropX2]
                #bg = cv2.cvtColor(bg, cv2.COLOR_BGR2HSV)

                img = self.filter(img, bg)
                cv2.imshow('dest', img)

                #refer_mask = np.ones((img.shape[0], img.shape[1]), dtype=img.dtype) * 255
                refer_mask = np.ones(img.shape, dtype=img.dtype) * 255
                fmask, smask, rmask = snapmodel.make_mask(img, refer_mask)

                cv2.imshow('fmask', fmask*127)
                cv2.imshow('smask', smask*127)
                cv2.imshow('rmask', rmask*127)
                cv2.waitKey(0)
                cv2.destroyAllWindows()

            self.croped.emit()

    def filter(self, img, bg):
        #x, y, w, h = 120, 75, 775, 370
        # 滤波
        bg = cv2.medianBlur(bg, 3)

        #bg = cv2.bilateralFilter(bg, 7, 75, 75)
        #bg = cv2.bilateralFilter(bg, 5, 30, 30)
        #corr = np.ones(bg.shape, dtype=bg.dtype) * np.min(bg)
        #corr = np.ones(bg.shape, dtype=bg.dtype)
        #bg = np.zeros(bg.shape, dtype=np.float) + bg
        #bg = bg * 1.01
        #bg = bg.astype(np.uint8) - corr
        #cv2.imshow('bg', bg)
        img = img * 1.01
        img = img.astype(np.uint8)
        cv2.imshow('img', img)
        #bg = np.subtract(bg, corr)
        #bg = bg + corr

        #imdst = np.subtract(bg, img) + corr
        #imdst = bg - img + corr
        imdst = bg - img
        return imdst

    def getCropImage(self):
        if self.cropImage is not None:
            img = self.cropImage.copy()
        if self.affine is not None:
            #cv2.warpPerspective -- 4点仿射
            return cv2.warpAffine(img, self.affine, (img.shape[1], img.shape[0])) 
        return img

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            if not self.isCrop:
                self.isCrop = True
                self.bgn = (event.x(), event.y())
        elif event.button() == Qt.RightButton:
            texts = label_eval.eval_text(label_eval.get_barcode(self.image))
            cv2.waitKey(0)
            cv2.destroyAllWindows()
            print(texts)

    def mouseMoveEvent(self, event):
        if self.isCrop:
            self.curr = (event.x(), event.y())
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton:
            if self.isCrop:
                self.isCrop = False
                self.end = (event.x(), event.y())
                self.crop()

    def paintEvent(self, event):
        super().paintEvent(event)
        if self.isCrop:
            qp = QPainter()
            qp.begin(self)
            self.drawRect(event, qp)
            qp.end()
        
    def drawRect(self, event, qp):
        qp.setPen(QColor(255,0,0))
        #print('{}, {}'.format(self.bgn, self.curr))
        x, y = self.bgn[0], self.bgn[1]
        w, h = self.curr[0] - x, self.curr[1] - y
        qp.drawRect(x, y, w, h)

    def resizeEvent(self, event):
        self.resized.emit()

class Widget(QWidget):
    def __init__(self, mode='text', path='.', file=''):
        super().__init__()
        self.mode = mode
        self.path = path
        self.file = os.path.join(path, file)
        self.setupUi()
        self.initImgList()
        self.btnPrev.clicked.connect(self.prev)
        self.btnNext.clicked.connect(self.next)
        #self.hsStretch.valueChanged.connect(self.setAffine)
        self.hsStretch.sliderReleased.connect(self.setAffine)
        self.lbImg.croped.connect(self.ocr)
        self.lbImg.resized.connect(self.view)
        self.leImgName.returnPressed.connect(self.view)
        self.image = None

    def setupUi(self):
        self.setObjectName("Form")
        self.resize(954, 650)
        self.gridLayout = QGridLayout(self)
        self.gridLayout.setObjectName("gridLayout")
        self.btnPrev = QPushButton(self)
        self.btnPrev.setObjectName("btnPrev")
        self.gridLayout.addWidget(self.btnPrev, 2, 2, 1, 1)
        self.btnNext = QPushButton(self)
        self.btnNext.setObjectName("btnNext")
        self.gridLayout.addWidget(self.btnNext, 2, 3, 1, 1)
        self.leImgName = QLineEdit(self)
        #self.leImgName.setReadOnly(True)
        self.leImgName.setObjectName("leImgName")
        self.gridLayout.addWidget(self.leImgName, 0, 2, 1, 2)
        self.btnRealTime = QRadioButton(self)
        self.btnRealTime.setMaximumSize(QSize(100, 16777215))
        self.btnRealTime.setObjectName("btnRealTime")
        self.gridLayout.addWidget(self.btnRealTime, 0, 0, 1, 1)
        self.btnLoadImage = QRadioButton(self)
        self.btnLoadImage.setMaximumSize(QSize(100, 16777215))
        self.btnLoadImage.setObjectName("btnLoadImage")
        self.gridLayout.addWidget(self.btnLoadImage, 0, 1, 1, 1)
        #self.lbImg = ImageLabel(self)
        #self.lbImg.setText("Image")
        #self.gridLayout.addWidget(self.lbImg, 1, 0, 1, 2)
        #self.lbImg.setAlignment(Qt.AlignCenter)
        self.lbImg = ImageLabel(self, self.mode)
        self.lbImg.setMinimumSize(QSize(622, 568))
        self.lbImg.setObjectName("lbImg")
        self.lbImg.setAlignment(Qt.AlignCenter)
        self.gridLayout.addWidget(self.lbImg, 1, 0, 1, 4)
        self.grpResult = QGroupBox(self)
        self.grpResult.setMaximumSize(QSize(260, 16777215))
        self.grpResult.setObjectName("grpResult")
        self.gridLayout_2 = QGridLayout(self.grpResult)
        self.gridLayout_2.setObjectName("gridLayout_2")
        self.hsStretch = QSlider(self.grpResult)
        self.hsStretch.setOrientation(Qt.Horizontal)
        self.hsStretch.setObjectName("hsStretch")
        self.hsStretch.setMinimum(-30)
        self.hsStretch.setMaximum(30)
        self.hsStretch.setValue(0)
        self.gridLayout_2.addWidget(self.hsStretch, 1, 0, 1, 2)
        self.lb0 = QLabel(self.grpResult)
        self.lb0.setObjectName("lb0")
        self.gridLayout_2.addWidget(self.lb0, 2, 0, 1, 1)
        self.lbTxt0 = QLabel(self.grpResult)
        self.lbTxt0.setObjectName("lbTxt0")
        self.lbTxt0.setMinimumSize(QSize(180, 0))
        self.gridLayout_2.addWidget(self.lbTxt0, 2, 1, 1, 1)
        self.lb90 = QLabel(self.grpResult)
        self.lb90.setObjectName("lb90")
        self.gridLayout_2.addWidget(self.lb90, 3, 0, 1, 1)
        self.lb270 = QLabel(self.grpResult)
        self.lb270.setObjectName("lb270")
        self.gridLayout_2.addWidget(self.lb270, 5, 0, 1, 1)
        self.lbTxt270 = QLabel(self.grpResult)
        self.lbTxt270.setObjectName("lbTxt270")
        self.gridLayout_2.addWidget(self.lbTxt270, 5, 1, 1, 1)
        self.lbTxt90 = QLabel(self.grpResult)
        self.lbTxt90.setObjectName("lbTxt90")
        self.gridLayout_2.addWidget(self.lbTxt90, 3, 1, 1, 1)
        self.lb180 = QLabel(self.grpResult)
        self.lb180.setObjectName("lb180")
        self.gridLayout_2.addWidget(self.lb180, 4, 0, 1, 1)
        self.lbTxt180 = QLabel(self.grpResult)
        self.lbTxt180.setObjectName("lbTxt180")
        self.gridLayout_2.addWidget(self.lbTxt180, 4, 1, 1, 1)
        self.lbStretch = QLabel(self.grpResult)
        self.lbStretch.setMaximumSize(QSize(16777215, 20))
        self.lbStretch.setObjectName("lbStretch")
        self.gridLayout_2.addWidget(self.lbStretch, 0, 0, 1, 1)
        self.gridLayout.addWidget(self.grpResult, 0, 4, 3, 1)
        self.btnSnap = QPushButton(self)
        self.btnSnap.setObjectName("btnSnap")
        self.gridLayout.addWidget(self.btnSnap, 2, 0, 1, 2)

        self.retranslateUi()
        QMetaObject.connectSlotsByName(self)

    def retranslateUi(self):
        _translate = QCoreApplication.translate
        self.setWindowTitle(_translate("Form", "Form"))
        self.btnSnap.setText(_translate("Form", "Snapshot"))
        self.btnPrev.setText(_translate("Form", "Prev"))
        self.btnNext.setText(_translate("Form", "Next"))
        self.grpResult.setTitle(_translate("Form", "Result"))
        self.btnRealTime.setText(_translate("Form", "RealTime"))
        self.btnLoadImage.setText(_translate("Form", "LoadImage"))
        self.lbStretch.setText(_translate("Form", "Stretch"))

    def initImgList(self):
        self.imgList = []
        try:
            self.imgList = [os.path.join(self.path, f) for f in os.listdir(self.path) if f[-4:].lower() == '.jpg']
            if len(self.imgList) > 0:
                if self.file in self.imgList:
                    self.currImgId = self.imgList.index(self.file)
                else:
                    self.currImgId = 0
                name = self.imgList[self.currImgId]
                self.leImgName.setText(name)
                self.view()
        except Exception as e:
            self.leImgName.setText(e)

    def view(self):
        try:
            name = self.leImgName.text()
            self.lbImg.showImage(name)
        except Exception as e:
            self.leImgName.setText(e)

    def prev(self):
        if self.currImgId > 0:
            self.currImgId -= 1
        else:
            self.currImgId = len(self.imgList) - 1
        name = self.imgList[self.currImgId]
        self.leImgName.setText(name)
        self.view()

    def next(self):
        if self.currImgId < len(self.imgList) - 1:
            self.currImgId += 1
        else:
            self.currImgId = 0
        name = self.imgList[self.currImgId]
        self.leImgName.setText(name)
        self.view()

    def ocr(self, reset=True):
        if reset:
            self.hsStretch.setValue(0)
            self.lbImg.setAffine(self.hsStretch.value())

        img = self.lbImg.getCropImage()

        # 锐化
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        img = ImageEnhance.Sharpness(img).enhance(5)
        img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

        # 锐化
        #kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], np.float32)
        #kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]], np.float32)
        #img = cv2.filter2D(img, -1, kernel=kernel)
        #img = cv2.convertScaleAbs(img)

        lbs = [self.lb0, self.lb90, self.lb180, self.lb270]
        txts = [self.lbTxt0, self.lbTxt90, self.lbTxt180, self.lbTxt270]
        for i in range(4):
            lb, txt = lbs[i], txts[i]
            tmp = img.copy()
            for j in range(i):
                tmp = cv2.transpose(tmp)
                tmp = cv2.flip(tmp, 1)

            result = ''
            if self.mode == 'zbar':
                import pyzbar.pyzbar as zbar
                content = zbar.decode(tmp)
                if len(content) > 0:
                    result = content[0]
            elif self.mode == 'dm':
                import pylibdmtx.pylibdmtx as dmtx
                content = dmtx.decode(tmp)
                if len(content) > 0:
                    result = content[0].data
            elif self.mode == 'text':
                import pytesseract
                result = pytesseract.image_to_string(tmp)
            else:
                print(f'{self.mode} is undefined.')
            txt.setText(f'{result}')

            tmp = QImage(tmp[:], tmp.shape[1], tmp.shape[0], tmp.shape[1] * 3, QImage.Format_RGB888)
            tmp = tmp.scaled(QSize(lb.width(), lb.height()), Qt.KeepAspectRatio)
            lb.setPixmap(QPixmap.fromImage(tmp))

    def setAffine(self):
        self.lbImg.setAffine(self.hsStretch.value())
        self.ocr(False)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--mode', default='text', help='Identify type: [text|zbar|dm|mask].')
    parser.add_argument('-p', '--path', default='.', help='Path of the picture that will be show.')
    parser.add_argument('-f', '--file', default='', help='Default picture that will be show.')
    args = parser.parse_args()
    mode = args.mode
    path = args.path
    file = args.file

    app = QApplication(sys.argv)
    w = Widget(mode=mode, path=path, file=file)
    w.show()
    sys.exit(app.exec())

